package state.demo1;

import org.apache.flink.api.common.state.CheckpointListener;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.OperatorStateStore;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

public class TestSource extends RichParallelSourceFunction<String>
    implements CheckpointedFunction, CheckpointListener {

  /**
   * operator-state为啥不支持valuestate呢？ 像这种一个值的存储
   */
  private final ListStateDescriptor<Long> beginIndexDesc = new ListStateDescriptor<Long>("beginIndex", Long.class);
  private transient ListState<Long> beginIndexState;
  private Map<Long, Long> pendingCommitOffsets = new HashMap<>();


  private boolean isCancel = false;
  private long currentIndex = 0;

  @Override
  public void open(Configuration parameters) {
  }

  @Override
  public void run(SourceContext<String> ctx) throws Exception {

    StringBuffer sb = new StringBuffer();
    for (int i = 0; i < 100000; i++) {
      sb.append("FileChannelManager removed spill file directoryFileChannelManager removed spill file directoryFileChannelManager removed spill file directory");
    }
    String msg = "MSG_" + sb.toString();


    System.out.println("7h_debug beginIndex = " + currentIndex);

    while (!isCancel) {

      String data = msg + currentIndex++;
      ctx.collect(data);

      Thread.sleep(1000);
    }
  }

  @Override
  public void cancel() {
    isCancel = true;
  }

  /**
   * 算子在启动时，恢复状态时调用
   */
  @Override
  public void initializeState(FunctionInitializationContext context) throws Exception {

    System.out.println("7h_debug initializeState call");
    OperatorStateStore operatorStateStore = context.getOperatorStateStore();
    this.beginIndexState = operatorStateStore.getListState(beginIndexDesc);

    if (context.isRestored()) {
      for (Long l : this.beginIndexState.get()) {
        this.currentIndex = l;
        break;
      }
    }
  }

  /**
   * snapshotState: 状态做持久化时调用。
   * 其实就算我们不写 snapshotState 方法也可以，Flink 会自动把上面的 ListState<String> 持久化，
   * snapshotState 给开发者实现特殊逻辑使用的，举例：在做 cp 时，可以从 ListState<String> l 删除一些不要的数据，添加一些特殊的数据。
   */
  @Override
  public void snapshotState(FunctionSnapshotContext context) throws Exception {
    System.out.println("7h_debug snapshotState call");

    pendingCommitOffsets.put(context.getCheckpointId(), currentIndex);

    beginIndexState.update(Collections.singletonList(currentIndex));

  }

  @Override
  public void notifyCheckpointComplete(long checkpointId) throws Exception {

    Long offset = pendingCommitOffsets.get(checkpointId);
    if (offset != null) {
      // commit offset
      System.out.println("7h_debug notifyCheckpointComplete call");
    }
  }
}
